package com.tydic.watermark;

import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.java.tuple.Tuple;
import org.apache.flink.streaming.api.TimeCharacteristic;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.AssignerWithPeriodicWatermarks;
import org.apache.flink.streaming.api.functions.windowing.WindowFunction;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.triggers.ContinuousEventTimeTrigger;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer010;
import org.apache.flink.util.Collector;

import javax.annotation.Nullable;
import java.text.SimpleDateFormat;
import java.util.*;

public class KafkaWatermarkTest {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setStreamTimeCharacteristic(TimeCharacteristic.EventTime);

        Properties properties = new Properties();
        properties.put("bootstrap.servers", "68.61.64.102:9092,68.61.64.103:9092");
        properties.put("group.id", "group01");

        properties.put("auto.offset.reset", "earliest");
        // properties.put("key.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
        // properties.put("value.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");

        DataStream<KafkaEvent> dataStream = env.addSource(new FlinkKafkaConsumer010<KafkaEvent>("test_topic", new KafkaEventSchema(), properties))
                .assignTimestampsAndWatermarks(new CustomWatermarkExtractor());

        dataStream.keyBy("xm").window(TumblingEventTimeWindows.of(Time.seconds(3)))
                .trigger(ContinuousEventTimeTrigger.of(Time.minutes(1))).apply(new WindowFunction<KafkaEvent, Object, Tuple, TimeWindow>() {
            @Override
            public void apply(Tuple tuple, TimeWindow timeWindow, Iterable<KafkaEvent> input, Collector<Object> out) throws Exception {
                SimpleDateFormat format = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS");
                List<KafkaEvent> events = new ArrayList<>();
                for (KafkaEvent kafkaEvent : input) {
                    events.add(kafkaEvent);
                }
                Collections.sort(events, new Comparator<KafkaEvent>() {
                    @Override
                    public int compare(KafkaEvent o1, KafkaEvent o2) {
                        if (o1.getCtime() > o2.getCtime()) {
                            return 1;
                        } else if (o1.getCtime() == o2.getCtime()) {
                            return 0;
                        } else {
                            return -1;
                        }
                    }
                });
                String result = tuple.getField(0) + "," + events.size() + ",First event:" + format.format(events.get(0).getCtime())
                        + ",Last event:" + format.format(events.get(events.size() - 1).getCtime())
                        + ",Window start:" + format.format(timeWindow.getStart()) + ",Window end:" + format.format(timeWindow.getEnd());
                out.collect(result);
            }
        }).print();

        dataStream.map(new RichMapFunction<KafkaEvent, Object>() {
            @Override
            public Object map(KafkaEvent kafkaEvent) throws Exception {
                return kafkaEvent;
            }
        });

        env.execute("Kafka watermark...");


    }

    private static class CustomWatermarkExtractor implements AssignerWithPeriodicWatermarks<KafkaEvent> {

        private SimpleDateFormat format = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS");
        private long currentMaxTimeStamp = 0L;
        private long maxAllowedLateness = 10000L; // 最大允许的乱序时间是10S
        private Watermark watermark;

        @Nullable
        @Override
        public Watermark getCurrentWatermark() {
            watermark = new Watermark(currentMaxTimeStamp - maxAllowedLateness);
            return watermark;
        }

        @Override
        public long extractTimestamp(KafkaEvent kafkaEvent, long previousElementTimestamp) {
            long timestamp = kafkaEvent.getCtime();
            currentMaxTimeStamp = Math.max(timestamp, currentMaxTimeStamp);
            System.err.println("Extract time:" + kafkaEvent.getCtime() + " "
                    + format.format(kafkaEvent.getCtime()) + " "
                    + "CurrentMaxTimestamp:" + format.format(currentMaxTimeStamp) + " "
                    + watermark.toString());
            return kafkaEvent.getCtime();
        }
    }
}
